﻿/* ===================================尊重原著======================================
 * 创 建 者：mao xiao jun 
 * 创建日期：2021-12-12
 * CLR Version : any
 * ===================================尊重原著=====================================*/
using System;
using System.Collections.Concurrent;
using System.IO;
using System.Linq;
using System.Net.Sockets;
using System.Security.Cryptography;
using System.Threading;
using System.Threading.Tasks;

namespace MTP
{
    abstract class MtpDispatch
    {
        protected static readonly byte[] PROTOCOL_HEADER = { 0x4D, 0x54, 0x50 };


        private readonly ConcurrentDictionary<Guid, AsyncBufferStream> concurrentDictionary = new ConcurrentDictionary<Guid, AsyncBufferStream>();

        /// <summary>
        /// 单包最大长度 最多2个字节表示长度
        /// </summary>
        private const int MAX_LENGTH = 1 << 16;


        protected readonly byte[] cipher;

        public MtpDispatch(string cipher)
        {
            this.cipher = MD5Sign(cipher);
        }

        protected async Task ReadNext(NetworkStream networkStream, CancellationToken cancellationToken)
        {
            byte[] bytes = new byte[22];
            int length = await networkStream.ReadAsync(bytes, 0, bytes.Length, cancellationToken);
            if (length != bytes.Length)
                throw new SocketException();
            Guid id = new Guid(bytes.AsSpan().Slice(0, 16));
            var packageType = (PackageType)bytes[16];
            switch (packageType)
            {
                case PackageType.Void:
                    ReadVoid(id);
                    break;
                case PackageType.Error:
                    Console.WriteLine("Error");
                    ReadError(id, new Exception());
                    break;
                default:
                    ushort bodyLength = BitConverter.ToUInt16(bytes, 17);
                    ushort sequence = BitConverter.ToUInt16(bytes, 19);
                    bool isEnd = BitConverter.ToBoolean(bytes, 21);
                    await ReadDataBefore(networkStream, id, sequence, isEnd, bodyLength);
                    break;
            }
        }


        private async Task ReadDataBefore(NetworkStream networkStream, Guid guid, ushort sequence, bool isEnd, int bodyLength)
        {
            byte[] bytes = new byte[bodyLength];
            int length = await networkStream.ReadAsync(bytes, 0, bytes.Length);
            if (length != bytes.Length)
                throw new SocketException();
            var buffer = concurrentDictionary.GetOrAdd(guid, guid => new AsyncBufferStream());
            string header = null;
            if (sequence == 0)
            {
                header = ReaderHeader(ref bytes);
            }
            if (buffer.Append(sequence, bytes, isEnd))
            {
                concurrentDictionary.TryRemove(guid, out _);
            }
            if (sequence == 0)
            {
                Task.Run(() => ReadData(guid, buffer, header, networkStream)).GetAwaiter();
            }
        }


        protected abstract Task ReadData(Guid guid, Stream stream, string header, NetworkStream networkStream);


        protected virtual void ReadError(Guid guid, Exception exception)
        {

        }

        protected virtual void ReadVoid(Guid guid)
        {

        }


        protected async Task SendVoid(Guid guid, NetworkStream networkStream)
        {
            var guidBytes = guid.ToByteArray();
            byte[] header = new byte[22];
            Array.Copy(guidBytes, header, 16);
            header[16] = (byte)PackageType.Void;
            await networkStream.WriteAsync(header, 0, header.Length);
        }

        protected async Task SendError(Guid guid, NetworkStream networkStream, Exception e)
        {
            var guidBytes = guid.ToByteArray();
            byte[] header = new byte[22];
            Array.Copy(guidBytes, header, 16);
            header[16] = (byte)PackageType.Error;


            byte[] bytes = new byte[MAX_LENGTH];

            byte[] lengthBytes = BitConverter.GetBytes((ushort)1);


            Array.Copy(lengthBytes, 0, bytes, 17, lengthBytes.Length);

            await networkStream.WriteAsync(bytes, 0, bytes.Length);
        }


        private byte[] GetHeaderBytes(string header)
        {
            if (string.IsNullOrEmpty(header))
            {
                return BitConverter.GetBytes((ushort)0);
            }
            byte[] bytes = System.Text.Encoding.UTF8.GetBytes(header);
            var data = BitConverter.GetBytes((ushort)bytes.Length);
            Array.Resize(ref data, data.Length + bytes.Length);
            Array.Copy(bytes, 0, data, 2, bytes.Length);
            return data;
        }

        private string ReaderHeader(ref byte[] bytes)
        {
            var lenght = BitConverter.ToUInt16(bytes);
            if (lenght == 0)
            {
                bytes = bytes.AsSpan(2 + lenght).ToArray();
                return string.Empty;
            }
            string header = System.Text.Encoding.UTF8.GetString(bytes, 2, lenght);
            bytes = bytes.AsSpan(2 + lenght).ToArray();
            return header;
        }


        protected async Task SendData(Guid guid, NetworkStream networkStream, Stream stream, string header)
        {
            var guidBytes = guid.ToByteArray();

            var headerBytes = GetHeaderBytes(header);
            ushort sequence = 0;
            int length;
            while (true)
            {
                byte[] bytes = new byte[MAX_LENGTH];

                length = await stream.ReadAsync(bytes, 0, bytes.Length - headerBytes.Length);

                if (length == 0 && sequence > 0) break;

                byte[] lengthBytes = BitConverter.GetBytes((ushort)(length + headerBytes.Length));

                byte[] seqBytes = BitConverter.GetBytes(sequence);

                byte[] data = new byte[22 + length + headerBytes.Length];

                Array.Copy(guidBytes, 0, data, 0, 16);

                // 包类型
                data[16] = (byte)PackageType.Data;

                Array.Copy(lengthBytes, 0, data, 17, lengthBytes.Length);

                Array.Copy(seqBytes, 0, data, 19, seqBytes.Length);

                // 是否结尾包
                data[21] = (byte)(length < MAX_LENGTH ? 1 : 0);

                if (headerBytes.Length > 0)
                    Array.Copy(headerBytes, 0, data, 22, headerBytes.Length);


                Array.Copy(bytes, 0, data, 22 + headerBytes.Length, length);


                await networkStream.WriteAsync(data, 0, data.Length);

                headerBytes = new byte[0];

                sequence++;
            };
        }


        protected async ValueTask<bool> VerifyHandshake(NetworkStream networkStream)
        {
            try
            {
                byte[] bytes = new byte[PROTOCOL_HEADER.Length];
                int length = await networkStream.ReadAsync(bytes, 0, bytes.Length, new CancellationTokenSource(2000).Token);
                if (length != bytes.Length)
                    return false;
                if (!PROTOCOL_HEADER.SequenceEqual(bytes))
                    return false;
                bytes = new byte[cipher.Length];
                length = await networkStream.ReadAsync(bytes, 0, bytes.Length);
                if (length != bytes.Length)
                    return false;
                if (!cipher.SequenceEqual(bytes))
                    return false;
                return true;
            }
            catch (Exception)
            {
                return false;
            }
        }

        protected async Task SendHandshake(NetworkStream networkStream)
        {
            byte[] bytes = new byte[PROTOCOL_HEADER.Length + cipher.Length];
            Array.Copy(PROTOCOL_HEADER, bytes, PROTOCOL_HEADER.Length);
            Array.Copy(cipher, 0, bytes, 3, cipher.Length);
            await networkStream.WriteAsync(bytes);
        }


        /// <summary>
        /// MD5签名
        /// </summary>
        /// <param name="str"></param>
        /// <returns></returns>
        private static byte[] MD5Sign(string str)
        {
            var data = System.Text.Encoding.UTF8.GetBytes(str);
            using (MD5 md5 = MD5.Create())
            {
                return md5.ComputeHash(data);
            }
        }
    }

    internal enum PackageType
    {
        /// <summary>
        /// 空包
        /// </summary>
        Void = 0,

        /// <summary>
        /// 数据包
        /// </summary>
        Data = 1,

        /// <summary>
        /// 异常包
        /// </summary>
        Error = 2

    }
}
